package com.example.module_publish.rtmp;

import static android.media.MediaCodec.CONFIGURE_FLAG_ENCODE;
import static android.media.MediaFormat.KEY_BIT_RATE;
import static android.media.MediaFormat.KEY_COLOR_FORMAT;
import static android.media.MediaFormat.KEY_FRAME_RATE;
import static android.media.MediaFormat.KEY_I_FRAME_INTERVAL;
import static android.media.MediaFormat.KEY_MAX_INPUT_SIZE;

import android.media.MediaCodec;
import android.media.MediaCodecInfo;
import android.media.MediaFormat;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.Date;
import java.util.concurrent.LinkedBlockingQueue;

/**
 * @author yaotianxue
 * @description
 * @date 2025/3/7
 */
public class MediaEncoder {
    private static final String TAG = "MediaEncoder";
    private Config mConfig;
    private LinkedBlockingQueue<byte[]> videoQueue;

    private MediaCodec vEncoder; //视频编码器
    private boolean videoEncoderLoop;
    private Thread videoEncoderThread;
    private long presentationTimeUs;
    private MediaCodec.BufferInfo vBufferInfo = new MediaCodec.BufferInfo();
    public static MediaEncoder newInstance(Config config) {
        return new MediaEncoder(config);
    }

    private MediaEncoder(Config config) {
        this.mConfig = config;
    }




    public void startVideoEncode() {
        videoEncoderThread = new Thread() {
            @Override
            public void run() {
                presentationTimeUs = System.currentTimeMillis() * 1000;
                vEncoder.start();
                while (videoEncoderLoop && !Thread.interrupted()) {
                    try {
                        byte[] data = videoQueue.take(); //待编码的数据
                        encodeVideoData(data);
                    } catch (InterruptedException e) {
                        e.printStackTrace();
                        break;
                    }
                }

            }
        };
        videoEncoderLoop = true;
        videoEncoderThread.start();
    }

    /**
     * 视频编码
     *
     * @param dstByte
     */
    private void encodeVideoData(byte[] dstByte) {
        ByteBuffer[] inputBuffers = vEncoder.getInputBuffers();
        ByteBuffer[] outputBuffers = vEncoder.getOutputBuffers();

        int inputBufferId = vEncoder.dequeueInputBuffer(-1);
        if (inputBufferId >= 0) {
            // fill inputBuffers[inputBufferId] with valid data
            ByteBuffer bb = inputBuffers[inputBufferId];
            bb.clear();
            bb.put(dstByte, 0, dstByte.length);
            long pts = new Date().getTime() * 1000 - presentationTimeUs;
            vEncoder.queueInputBuffer(inputBufferId, 0, dstByte.length, pts, 0);
        }

        int outputBufferId = vEncoder.dequeueOutputBuffer(vBufferInfo, 0);
        if (outputBufferId >= 0) {
            // outputBuffers[outputBufferId] is ready to be processed or rendered.
            ByteBuffer bb = outputBuffers[outputBufferId];
            if (null != mCallback) {
                mCallback.outputVideoData(bb, vBufferInfo);
            }
            MediaFormat outputFormat = vEncoder.getOutputFormat();

            vEncoder.releaseOutputBuffer(outputBufferId, false);
        }

    }

    /**
     * 初始化视频编码器。
     *
     * @param width  视频的宽
     * @param height 视频的高
     * @throws IOException 创建编码器失败
     */
    public void initVideoEncoder(int width, int height) throws IOException {
        MediaFormat format = MediaFormat.createVideoFormat(MediaFormat.MIMETYPE_VIDEO_AVC,
                width, height);
        format.setInteger(KEY_MAX_INPUT_SIZE, 0);
        format.setInteger(KEY_BIT_RATE, mConfig.bitrate);
        format.setInteger(KEY_COLOR_FORMAT, MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420Planar);
        format.setInteger(KEY_FRAME_RATE,  mConfig.fps);
        format.setInteger(KEY_I_FRAME_INTERVAL, 1);
        vEncoder = MediaCodec.createEncoderByType(MediaFormat.MIMETYPE_VIDEO_AVC);
        vEncoder.configure(format, null, null, CONFIGURE_FLAG_ENCODE);
        vEncoder.setCallback(callback); //设置回调
        videoQueue = new LinkedBlockingQueue<>();
    }

    /**
     * 添加视频数据
     *
     * @param data
     */
    public void putVideoData(byte[] data) {
        try {
            videoQueue.put(data);
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
    }

    /**
     * stop视频编码
     */
    public void stopVideoEncode() {
        videoEncoderLoop = false;
        if(videoEncoderThread != null){
            videoEncoderThread.interrupt();
        }
        if(vEncoder != null){
            vEncoder.stop();
        }
    }


    /**
     * 释放视频编码器
     */
    public void releaseVideoEncoder() {
        if(vEncoder!=null){
            vEncoder.release();
        }
    }
    private Callback mCallback;

    public interface Callback {
        void outputVideoData(ByteBuffer bb, MediaCodec.BufferInfo info);

        void outputAudioData(ByteBuffer bb, MediaCodec.BufferInfo info);
    }
    /**
     * 设置回调
     *
     * @param callback 回调
     */
    public void setCallback(Callback callback) {
        this.mCallback = callback;
    }


    private MediaCodec.Callback callback;
    public void setCallback(MediaCodec.Callback callback) {
        this.callback = callback;
    }
}
